import numpy as np
import torch
from skimage.transform import rescale

from defense.base_defense import BaseDefense
# 不用retrain但需要transform input
class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('RT', model_name, config, retrain = False)

    def transform_input(self, images):
        size = images.shape[-1]
        resize = self.config['resize']
        # 数据结构转换
        if torch.is_tensor(images) is True:
            images = images.numpy()
        # 转换图像通道
        images = np.transpose(images, (0, 2, 3, 1))

        trans_images = []
        for image in images:
            # 1 随机调整层
            # 指定图像重新缩放到随机大小
            rnd = np.random.randint(size, resize)
            scale = (rnd * 1.0) / size
            r_image = rescale(image, scale, multichannel=True, preserve_range=True, mode='constant',
                                     anti_aliasing=False)
            # 2 随机填充层
            h = resize - rnd
            w = resize - rnd
            pad_l = np.random.randint(0, w)
            pad_r = w - pad_l
            pad_t = np.random.randint(0, h)
            pad_b = h - pad_t
            # 使用灰色像素填充图像至新的大小
            padded_image = np.pad(r_image, ((pad_t, pad_b), (pad_l, pad_r), (0, 0)), 'constant',
                                  constant_values=0.5)
            trans_images.append(padded_image)

        # 重新设置通道位置并转换数据结构
        trans_images = np.array(trans_images)
        trans_images = torch.from_numpy(np.transpose(trans_images, (0, 3, 1, 2))).float()

        return trans_images





